# -*- coding: utf-8 -*-
"""
Module Description: 用线程池作为异步任务池
"""
from Queue import Queue
import threading
import time
from framework.firefly_dc import runtime_logger as log
from framework.firefly_dc.context import action_context_storage
from utils import traceback_utils

# 关闭线程池的load
StopWorker = object()


class ThreadPool(object):
    """
    工作线程池,用于处理异步任务
    """
    # 认为是过长的任务执行时间
    TOO_LONG_EXEC_TIME = 1000
    # 认为是过长的任务队列
    TOO_LONG_TASK_QUEUE = 100

    def __init__(self, maxthreads=10, name=None):
        # 任务队列
        self.task_queue = Queue(0)
        # 工作中线程数量
        self.worker_num = 0
        # 线程池大小
        self.max = maxthreads
        # 线程池名称
        self.name = name or id(self)
        # 工作线程
        self.workers = []

    def start(self):
        """启动线程池"""
        while self.worker_num < self.max:
            self._start_worker()

    def _start_worker(self):
        """启动工作线程"""
        self.worker_num += 1
        name = 'PoolThread-{poll_name}_{tid}'.format(poll_name=self.name, tid=self.worker_num)
        thread = threading.Thread(target=self._hardworking, name=name)
        thread.start()
        self.workers.append(thread)

    def submit(self, func, *args, **kwargs):
        """
        提交新任务
        :param func: 执行任务的函数主体
        :param args: 函数参数
        :param kwargs: 函数参数
        """
        self.submit_with_callback(None, func, *args, **kwargs)

    def submit_with_callback(self, callback, func, *args, **kwargs):
        """
        提交含回调函数的新任务

        回调函数的参数有两个:
        - 第一个是任务函数主体执行是否正常完成
        - 第二个是任务函数主体的返回值

        如果任务函数主体执行时抛出异常,则返回值为None

        :param callback: 回调函数
        :param func: 执行任务的函数主体
        :param args: 函数参数
        :param kwargs: 函数参数
        :return:
        """
        o = (func, args, kwargs, callback, action_context_storage.ctx)
        self.task_queue.put(o)

        # 记录过长的任务队列日志
        task_queue_size = self.task_queue.qsize()
        if task_queue_size > self.TOO_LONG_TASK_QUEUE:
            log.warning('ThreadPool task queue too long: {task_len}'.format(task_len=task_queue_size))

    def stop(self):
        """关闭线程池"""
        while self.worker_num:
            self.task_queue.put(StopWorker)
            self.worker_num -= 1

        log.info('ThreadPool {name} stop.'.format(name=self.name))

    @staticmethod
    def _run_with_ctx(ctx, func, *args, **kwargs):
        """
        在上下文中执行方法
        :param ctx: 上下文对象
        :param func: 执行体
        :return:
        """
        if ctx is not None:
            with ctx:
                return func(*args, **kwargs)
        else:
            return func(*args, **kwargs)

    def _hardworking(self):
        """线程主体"""
        cur_thread = threading.currentThread()
        cur_thread_name = cur_thread.getName()
        log.info('WorkerThread {name} start hardworking.'.format(name=cur_thread_name))

        o = self.task_queue.get()

        while o is not StopWorker:
            func, args, kwargs, callback, ctx = o
            del o
            log.info('WorkerThread-({name}) run task: {task_name}'.format(
                name=cur_thread_name,
                task_name=func.__name__
            ))

            st = time.time()
            try:
                result = self._run_with_ctx(ctx, func, *args, **kwargs)
                successed = True
            except Exception as e:
                traceback_utils.print_exc(e)

                result = None
                successed = False
            finally:
                # 记录耗时过长的日志
                cost_time = (time.time() - st) * 1000
                if cost_time >= self.TOO_LONG_EXEC_TIME:
                    self._run_with_ctx(
                        ctx,
                        log.warning,
                        'WorkerThread-{name} run task: {task_name} cost time: {cost_time} ms.Ctx({ctx})'.format(
                            name=cur_thread_name,
                            task_name=func.__name__,
                            cost_time=cost_time,
                            ctx=ctx)
                    )

            if callback is not None:
                self._run_with_ctx(ctx, callback, *(successed, result))

            # wait for next load
            o = self.task_queue.get()

        log.info('Worker Thread {name} rest.'.format(name=cur_thread.getName()))


if __name__ == '__main__':
    def task():
        print 'sleep...'
        import requests
        import random
        response = requests.get(random.choice(['http://www.google.com', 'http://www.baidu.com']))
        print response
        print 'wake up.'

        return response.status_code

    def _callback(successed, result):
        if successed:
            print 'very good!'
        print 'cb:', result

    pool = ThreadPool(name='Test')
    pool.start()

    for i in range(100):
        if i % 2 == 0:
            pool.submit(task)
        else:
            pool.submit_with_callback(_callback, task)

    pool.submit(pool.stop)
